import subprocess
import threading
import time
import os
import signal
import logging
import psutil
from typing import Dict, List, Optional, Union

logger = logging.getLogger(__name__)

class OllamaProcessManager:
    """Ollama进程管理器，确保服务正常运行并提供进程隔离的调用方式"""
    
    def __init__(self):
        self.active_processes = {}
        self.lock = threading.Lock()
    
    def ensure_service_running(self) -> bool:
        """确保Ollama服务正在运行"""
        try:
            # 检查Ollama服务是否在运行
            ollama_running = False
            for proc in psutil.process_iter(['pid', 'name']):
                if 'ollama' in proc.info['name']:
                    ollama_running = True
                    break
                    
            if not ollama_running:
                logger.info("Ollama服务未运行，正在启动...")
                # 启动ollama服务
                subprocess.Popen(
                    ["ollama", "serve"],
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.DEVNULL
                )
                # 等待服务启动
                time.sleep(5)
                return True
                
            # 检查服务是否响应
            status_check = subprocess.run(
                ["curl", "-s", "http://localhost:11434/api/tags"],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                timeout=5
            )
            
            if status_check.returncode != 0:
                logger.warning("Ollama服务未响应，正在重启...")
                self.restart_service()
                return True
                
            return True
        except Exception as e:
            logger.error(f"确认Ollama服务状态时出错: {str(e)}")
            return False
    
    def restart_service(self) -> bool:
        """重启Ollama服务"""
        try:
            logger.info("正在重启Ollama服务...")
            
            # 终止所有活跃的Ollama进程
            self.terminate_all_processes()
            
            # 杀死所有现有的ollama进程
            for proc in psutil.process_iter(['pid', 'name']):
                if 'ollama' in proc.info['name']:
                    try:
                        process = psutil.Process(proc.info['pid'])
                        process.terminate()
                    except:
                        pass
            
            # 等待进程终止
            time.sleep(2)
            
            # 检查是否还有残留进程
            for proc in psutil.process_iter(['pid', 'name']):
                if 'ollama' in proc.info['name']:
                    try:
                        process = psutil.Process(proc.info['pid'])
                        process.kill()
                    except:
                        pass
            
            # 确保服务完全停止
            time.sleep(1)
            
            # 启动ollama服务
            subprocess.Popen(
                ["ollama", "serve"],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL
            )
            
            logger.info("已重启Ollama服务，等待服务就绪...")
            time.sleep(5)  # 等待服务启动
            return True
        except Exception as e:
            logger.error(f"重启Ollama服务失败: {str(e)}")
            return False
    
    def run_with_timeout(self, model: str, prompt: str, timeout: int = 300) -> Dict[str, any]:
        """
        使用独立进程运行Ollama模型，具有超时保护
        
        参数:
            model: 模型名称
            prompt: 提示词
            timeout: 超时时间(秒)
            
        返回:
            结果字典: {"success": bool, "text": str, "error": str}
        """
        # 确保服务正在运行
        if not self.ensure_service_running():
            return {
                "success": False,
                "text": "",
                "error": "无法确保Ollama服务正在运行"
            }
        
        result = {
            "success": False,
            "text": "",
            "error": "",
            "process_id": None
        }
        
        # 使用线程来设置超时
        completion_event = threading.Event()
        
        # 创建临时文件存储提示词
        prompt_file = f"prompt_{int(time.time())}.txt"
        
        try:
            # 写入提示词到临时文件
            with open(prompt_file, 'w', encoding='utf-8') as f:
                f.write(prompt)
            
            # 构建命令
            cmd = f'ollama run {model} -f {prompt_file}'
            
            # 进程ID锁
            with self.lock:
                process_id = len(self.active_processes) + 1
                result["process_id"] = process_id
            
            def run_process():
                try:
                    # 启动进程
                    process = subprocess.Popen(
                        cmd,
                        shell=True,
                        stdout=subprocess.PIPE,
                        stderr=subprocess.PIPE,
                        text=True,
                        preexec_fn=os.setsid  # 创建新进程组，便于终止
                    )
                    
                    # 记录活跃进程
                    with self.lock:
                        self.active_processes[process_id] = {
                            "process": process,
                            "start_time": time.time(),
                            "command": cmd
                        }
                    
                    # 获取输出
                    stdout, stderr = process.communicate()
                    
                    # 更新结果
                    if process.returncode == 0:
                        result["success"] = True
                        result["text"] = stdout.strip()
                    else:
                        result["error"] = stderr.strip()
                        
                except Exception as e:
                    result["error"] = str(e)
                finally:
                    # 移除进程记录
                    with self.lock:
                        if process_id in self.active_processes:
                            del self.active_processes[process_id]
                    
                    # 通知主线程
                    completion_event.set()
            
            # 启动处理线程
            thread = threading.Thread(target=run_process)
            thread.daemon = True  # 允许程序退出时终止线程
            thread.start()
            
            # 等待完成或超时
            if not completion_event.wait(timeout=timeout):
                # 超时 - 终止进程
                with self.lock:
                    if process_id in self.active_processes:
                        process_info = self.active_processes[process_id]
                        process = process_info["process"]
                        
                        try:
                            # 终止整个进程组
                            os.killpg(os.getpgid(process.pid), signal.SIGTERM)
                        except:
                            # 备用方法
                            try:
                                process.terminate()
                                process.kill()
                            except:
                                pass
                        
                        del self.active_processes[process_id]
                
                result["error"] = "生成超时，已终止进程"
                logger.warning(f"进程 {process_id} 执行超时，已终止")
            
            return result
            
        except Exception as e:
            logger.error(f"执行进程时出错: {str(e)}")
            return {
                "success": False,
                "text": "",
                "error": str(e)
            }
        finally:
            # 清理临时文件
            if os.path.exists(prompt_file):
                try:
                    os.remove(prompt_file)
                except:
                    pass
    
    def terminate_process(self, process_id: int) -> bool:
        """终止指定的进程"""
        with self.lock:
            if process_id in self.active_processes:
                process_info = self.active_processes[process_id]
                process = process_info["process"]
                
                try:
                    # 终止整个进程组
                    os.killpg(os.getpgid(process.pid), signal.SIGTERM)
                except:
                    # 备用方法
                    try:
                        process.terminate()
                        process.kill()
                    except:
                        pass
                
                # 更新状态
                del self.active_processes[process_id]
                logger.info(f"已终止进程 {process_id}")
                return True
            else:
                logger.warning(f"进程 {process_id} 不存在或已终止")
                return False
    
    def terminate_all_processes(self) -> int:
        """终止所有活跃进程，返回终止的进程数量"""
        terminated_count = 0
        
        with self.lock:
            process_ids = list(self.active_processes.keys())
            
            for pid in process_ids:
                if self.terminate_process(pid):
                    terminated_count += 1
        
        return terminated_count
    
    def get_active_processes(self) -> Dict[int, Dict]:
        """获取所有活跃进程的信息"""
        with self.lock:
            result = {}
            current_time = time.time()
            
            for pid, info in self.active_processes.items():
                result[pid] = {
                    "start_time": info["start_time"],
                    "running_time": current_time - info["start_time"],
                    "command": info["command"]
                }
            
            return result
    
    def check_resources(self) -> Dict[str, any]:
        """检查系统资源状态"""
        try:
            result = {
                "memory": {
                    "available": "Unknown",
                    "total": "Unknown",
                    "percent": "Unknown"
                },
                "gpu": []
            }
            
            # 检查系统内存
            mem_info = {}
            with open('/proc/meminfo', 'r') as f:
                for line in f:
                    parts = line.split(':')
                    if len(parts) >= 2:
                        key = parts[0].strip()
                        value = parts[1].strip()
                        if 'kB' in value:
                            value = int(value.replace('kB', '').strip()) * 1024
                        mem_info[key] = value
            
            total = int(mem_info.get('MemTotal', 0))
            available = int(mem_info.get('MemAvailable', 0))
            
            if total > 0:
                used_percent = (total - available) / total * 100
                result["memory"] = {
                    "available": f"{available / (1024 * 1024 * 1024):.2f} GB",
                    "total": f"{total / (1024 * 1024 * 1024):.2f} GB",
                    "percent": f"{used_percent:.1f}%",
                    "high_usage": used_percent > 90
                }
            
            # 检查GPU显存(如果有)
            try:
                nvidia_smi = subprocess.check_output(['nvidia-smi', '--query-gpu=memory.used,memory.total', '--format=csv,nounits,noheader'])
                values = nvidia_smi.decode('ascii').strip().split('\n')
                
                for i, line in enumerate(values):
                    mem_used, mem_total = map(int, line.split(','))
                    mem_percent = mem_used / mem_total * 100
                    
                    result["gpu"].append({
                        "id": i,
                        "memory_used": f"{mem_used / 1024:.2f} GB",
                        "memory_total": f"{mem_total / 1024:.2f} GB",
                        "percent": f"{mem_percent:.1f}%",
                        "high_usage": mem_percent > 90
                    })
            except:
                pass  # 没有GPU或无法查询
                
            return result
        except Exception as e:
            logger.warning(f"资源检查失败: {str(e)}")
            return {"error": str(e)}
